import copy
from typing import List

from torch.distributed._tensor.placement_types import (
    _Partial,
    DTensorSpec,
    Replicate,
    Shard,
    TensorMeta,
)
from ..shard.placement_types import PlacementStrategy

from .strategy_generator import StrategyGenerator
__all__ = ['TensorConstructorGenerator']


class TensorConstructorGenerator(StrategyGenerator):
    """
    TensorConstructorGenerator which deals with
    the sharding strategies for tensor constructor operation, such as torch.arange.
    """

    # def validate(self) -> bool:
    #     return super().validate()

    def collate_strategies(self) -> List[PlacementStrategy]:
        strategy_list = []
        dim_partition_dict_mapping = {
            "output": {},
        }
        # communication_action_mapping = {}
        # sharding_spec_mapping = self.to_sharding_spec_mapping(dim_partition_dict_mapping)
        # 直接复制
        name = 'Replica Tensor Constructor'
        out_spec = DTensorSpec(self.device_mesh, (Replicate(),), self.op_data['output'].data)
        sharding_specs_mapping = {'output':out_spec}
        sharding_specs_data = self.get_sharding_specs(sharding_specs_mapping)
        # communication TODO
        strategy = ShardingStrategy(name=name,sharding_specs=sharding_specs_data)
        strategy_list.append(strategy)

        return strategy_list
